##Correct paths below - if you're keeping with the format in the zip-file, just have R use the root folder as its working directory, and hopefully everything will work.

##this is the code that implements RJMCMC - and probably where you can offer the most useful improvements:
source("./changepoints/cp.R")

##formatting the data...
track.files <- list.files("./data/newspaper", full.names=TRUE)
tracks <- list()
for (i in (1:length(track.files))){
  v<-as.matrix(read.table(track.files[i]))
  valid <- which(v[,4]>0)
  tracks[[i]]<-cbind(5/168*(valid-1), v[valid,1], v[valid,3])
}

##this is a useful to understand what's going on. Just use plot.tracks() to plot everyone and look up person id from the colour in the plot. E.g. "black" is 1, "red" is 2.
plot.tracks <- function(start.time=0, end.time=5, cols=c("black", "red", "green", "blue", "orange", "pink", "yellow", "cyan")){
  if (length(tracks) > length(cols)){stop("Specify more colours")}
  xlim<-range(lapply(tracks, function(track){c(min(track[,2]), max(track[,2]))}))
  ylim<-range(lapply(tracks, function(track){c(min(-track[,3]), max(-track[,3]))}))
  plottable <- tracks[[1]][(tracks[[1]][,1] < end.time) & (tracks[[1]][,1] >=start.time), c(2,3)]
  plot(plottable[,1], -plottable[,2], xlim=xlim, ylim=ylim, type="l", col=cols[1])
  for (i in (1:length(tracks))){
    plottable <- tracks[[i]][(tracks[[i]][,1] < end.time) & (tracks[[i]][,1] >=start.time), c(2,3)]
    lines(plottable[,1], -plottable[,2], col=cols[i])
  }
}

##load("./data/pnewspaper.Rout")
##pdf(width=3, height=3*(33-15)/(3.5+2.5))
plot.tracks.beautiful <- function(){
  xlim <- c(-2.5, 3.5)
  #xlim <- c(-9, 9)
  ylim <- c(15, 33)
  track <- tracks[[1]]
  par(mar=c(0,0,0,0))
  plot(c(), xlim=xlim, ylim=ylim, axes=FALSE, xlab=NULL, ylab=NULL)
  index <- c()
  for (i in 1:(length(tracks)-1)){
    for (j in (i+1):length(tracks)){
      index <- c(index, paste(i, j))
    }      
  }
  for (comb in 1:length(p)){
    actors <- as.numeric(unlist(strsplit(index[comb], " ")))
    track1 <- tracks[[actors[1]]]; track2 <- tracks[[actors[2]]];
    start1 <- min(which(track1[,1]>=track2[1,1]))-1
    start2 <- min(which(track2[1,1]<=track2[,1]))-1
    for (f in 1:dim(p[[comb]])[2]){
      if (p[[comb]][2,f]>0.5){        
        lines(c(track1[start1+f,2], track2[start2+f,2]), c(-track1[start1+f,3], -track2[start2+f,3]), col=rgb(1,0,0,0.8))
      }
      ## else {
      ##   #lines(c(track1[start1+f,2], track2[start2+f,2]), c(-track1[start1+f,3], -track2[start2+f,3]), col=rgb(1-2*p[[comb]][2,f], 1-2*p[[comb]][2,f], 1-2*p[[comb]][2,f], 0.1))
      ##   lines(c(track1[start1+f,2], track2[start2+f,2]), c(-track1[start1+f,3], -track2[start2+f,3]), col=rgb(0,0,1-2*p[[comb]][2,f], 0.1))
      ## }
      else if (p[[comb]][2,f]>0.05){
        lines(c(track1[start1+f,2], track2[start2+f,2]), c(-track1[start1+f,3], -track2[start2+f,3]), col=rgb(0,0,1-2*p[[comb]][2,f], 0.3))
      }
    }
  }  
  for (i in 1:length(tracks)){
    track <- tracks[[i]]
    lines(track[,2], -track[,3], xlim=xlim, ylim=ylim, type="l", lwd=3)
    points(track[1,2], -track[1,3], pch=16)
  }
}

plot.s.beautiful <- function(){
  D <- mkD(1,2)
  dist<- sqrt(D[2,]^2+D[3,]^2)
  ylim <- c(0,3.5)
  par(mar=c(4,4,0,0))
  plot(D[1,], dist, ylim=ylim, ylab="", xlab="Time", axes=FALSE)
  axis(2, at=0:3)
  axis(1, at=0:5)
  box()
  lines(p[[1]][1,], p[[1]][2,], col="black")
  abline(h=0.5, lty=2)  
}

mkD <- function(id1=1, id2=2, start.time=NULL, end.time=NULL){
  ##Gets data for two given Ids. Basically, just use e.g. D<-mkD(1,2) and you get all the data for when 1 *and* 2 are in view.
  if (is.null(start.time)){start.time=0; end.time=5}
  track1<- tracks[[id1]]; track2<-tracks[[id2]]
  start <- max(c(start.time, track1[1,1], track2[1,1]))
  end <- min(end.time, track1[length(track1[,1]),1],track2[length(track2[,1]),1])
  rbind(track1[(track1[,1]>=start&track1[,1]<=end),1], track1[(track1[,1]>=start&track1[,1]<=end),2] - track2[(track2[,1]>=start&track2[,1]<=end),2], track1[(track1[,1]>=start&track1[,1]<=end),3] - track2[(track2[,1]>=start&track2[,1]<=end),3])
}

doforall <- function(){
  p <- list()
  for (i in (1:5)){
    for (j in (i+1):6){
      cat(i, j, "\n")
      D <- mkD(i,j)
      set.seed(0)
      s<-getsample(D, N=5000, lambda=1)
      p[[length(p)+1]]<-rbind(D[1,], prob.regime(D[1,], s, 1))      
    }
  }
  p
}

prob.regime <- function(times, s, regime){
  ##calculates the posterior probability, at each time-point in times, that the two actors are "together" based on all the sampled changepoints.
  counts<-rep(0, length(times))
  for (i in 1:length(s)){
    for (j in 1:length(times)){
      counts[j] = counts[j] + (s[[i]]$ms[max(which(c(0, s[[i]]$cps)<=times[j]))]==regime)
    }
  }
  counts/length(s)
}

plot.s <- function(s, D){
  ##plots the change-points density esimate (red), the probability that two actors are together other time (green), the distance difference (circles) and the (scaled to fit in the plot) total distance (black line). Everything basically.
  dist<- sqrt(D[2,]^2+D[3,]^2)
  distdiff <- diff(dist)
  dist <- dist/max(dist) * max(distdiff)  
  plot(D[1,1:length(distdiff)], distdiff, ylim=c(min(distdiff), max(1, max(distdiff))), xlab="Time", ylab="")
  lines(D[1,], dist)
  cps <- get.cps(s)
  lines(density(cps), col="red")
  t<-D[1,]
  lines(t, prob.regime(t, s, 1), col="green")
  abline(h=0.5)
}


D <- mkD(3,5) ## put any actor indices you want here - 1 and 2 are the woman and man at the front of the video. Use plot.tracks if you want to know who is who. (The colour number is the actor index).
##plot.tracks()
set.seed(0) ## I like to do this, for reproducibility.
s<-getsample(D, N=5000, lambda=1)
psum(s)
#plot.s(s,D)##all the computation time is taken here - the actual sampling is comparatively very fast. Eventually we'll have to look at prob.regimes and make it a bit faster, it's a bit frustrating to have wait like this all the time, just to see the plot.


